# coding:utf-8
import mmap
import socket
import sys
import json
import struct

if sys.version_info[0] == 2:
    reload(sys)
    sys.setdefaultencoding('utf8')
    int_from_byte = ord


    def uint_from_bytes(b):
        if b:
            return int(b.encode("hex"), 16)
        return 0


    def int_to_ipv4(int_ip):
        packed_ip = struct.pack('!I', int_ip)
        return socket.inet_ntoa(packed_ip)


    def ipv4_to_int(ip):
        packed_ip = socket.inet_aton(ip)
        return int(packed_ip.encode('hex'), 16)


    def ipv6_to_int(ipv6_address):
        ipv6_bytes = socket.inet_pton(socket.AF_INET6, ipv6_address)
        return int(struct.unpack('!2Q', ipv6_bytes)[0] << 64 | struct.unpack('!2Q', ipv6_bytes)[1])
else:
    int_from_byte = lambda x: x
    uint_from_bytes = lambda x: int.from_bytes(x, 'big')


    def int_to_ipv4(ip_integer):
        return socket.inet_ntoa(ip_integer.to_bytes(4, byteorder='big'))


    def ipv4_to_int(ip):
        return int.from_bytes(socket.inet_aton(ip), byteorder='big')


    def ipv6_to_int(ipv6_address):
        return int.from_bytes(socket.inet_pton(socket.AF_INET6, ipv6_address), byteorder='big')


def is_ipv4(ip):
    '''
    判断ipv4
    '''
    try:
        socket.inet_pton(socket.AF_INET, ip)
    except AttributeError:
        try:
            socket.inet_aton(ip)
        except socket.error:
            return False
        return ip.count('.') == 3
    except socket.error:
        return False
    return True


def is_ipv6(ip):
    '''
    判断IPv6
    '''
    try:
        socket.inet_pton(socket.AF_INET6, ip)
    except socket.error:
        return False
    return True


def ipv4_map_ipv6(ipv4_address):
    """将IPv4地址映射到IPv6地址"""
    ipv4_bytes = socket.inet_pton(socket.AF_INET, ipv4_address)
    ipv6_bytes = b"\x00" * 10 + b"\xff" * 2 + ipv4_bytes
    ipv6_address = socket.inet_ntop(socket.AF_INET6, ipv6_bytes)
    return ipv6_address


def map_key_value(keys, values):
    '''
    key和value映射
    :param keys:
    :param values:
    :return:
    '''
    if len(keys) == len(values):
        count = len(keys)
        result_dict = dict(zip(keys[:count], values[:count]))
    else:
        count = len(values) - 1
        forepart_dict = dict(zip(keys[:count], values[:count]))
        keys_list = keys[-1]
        values_list = values[-1]
        multiAreas_name = keys[-2]
        end_dict = {multiAreas_name: [dict(zip(keys_list, row)) for row in values_list]}
        result_dict = dict(forepart_dict, **end_dict)
    return result_dict


class AWDB(object):
    def __init__(self, file_name):
        with open(file_name, 'rb') as db_file:
            self.byte_data = mmap.mmap(db_file.fileno(), 0, access=mmap.ACCESS_READ)
        self.metadata_len = uint_from_bytes(self.byte_data[:2])
        self.metadata = self.byte_data[2:2 + self.metadata_len]
        self.metadata = json.loads(self.metadata)
        self.node_count = self.metadata['node_count']
        self.ip_version = self.metadata['ip_version']
        self.decode_type = self.metadata['decode_type']
        self.byte_len = self.metadata['byte_len']
        self.columns = self.metadata['columns']
        self.start_len = 2 + self.metadata_len
        self.base_offset = self.start_len + self.node_count * self.byte_len * 2

    def find_tree_index(self, ip_address):
        if self.ip_version == '4_6':
            if type(ip_address) == int:
                if ip_address < 2 ** 32:
                    ip_address = int_to_ipv4(ip_address)
                    ip_address = ipv4_map_ipv6(ip_address)
            elif is_ipv4(ip_address):  # 处理这种情况 192.168.1.1
                ip_address = ipv4_map_ipv6(ip_address)

        if isinstance(ip_address, int):
            if self.ip_version == '4':
                if ip_address >= 2 ** 32:
                    raise ValueError(
                        'The database is IPv4 library, but you are using ipv6 queries! The IP is:{}'.format(ip_address))
                ip_count = 32
                int_ip = ip_address
            else:
                if ip_address < 2 ** 32:
                    raise ValueError(
                        'The database is IPv6 library, but you are using ipv4 queries! The IP is:{}'.format(ip_address))
                ip_count = 128
                int_ip = ip_address
        else:
            if is_ipv4(ip_address):
                if self.ip_version == '4':
                    int_ip = ipv4_to_int(ip_address)
                    ip_count = 32
                else:
                    raise ValueError(
                        'The database is IPv6 library, but you are using ipv4 queries! The IP is:{}'.format(ip_address))
            elif is_ipv6(ip_address):
                if self.ip_version == '6' or self.ip_version == '4_6':
                    int_ip = ipv6_to_int(ip_address)
                    ip_count = 128
                else:
                    raise ValueError('The database is IPv4 library, but you are using ipv6 queries! The IP is:{}',
                                     ip_address)
            else:
                raise ValueError('Invalid IP address {}!'.format(ip_address))

        node_index = 0
        if self.ip_version == '4_6':
            if (int(int_ip) >> 32) == 0xFFFF:
                node_index = 96

        for i in range(ip_count - node_index - 1, -1, -1):
            if node_index >= self.node_count:
                break
            bit_value = (int_ip >> i) & 1
            offset = node_index * self.byte_len * 2 + bit_value * self.byte_len + self.start_len
            node_byte = self.byte_data[offset:offset + self.byte_len]
            node_index = uint_from_bytes(node_byte)

        if node_index == self.node_count:
            print('Invalid node_index in search tree')
            return ''
        if node_index > self.node_count:
            return node_index
        raise ValueError('Invalid node_index in search tree')

    def decode_arrary(self, new_offset):
        data_len = int_from_byte(self.byte_data[new_offset])
        new_offset += 1
        array = []
        for _ in range(data_len):
            value, new_offset = self.decode_content(new_offset)
            array.append(value)
        return array, new_offset

    def decode_str(self, new_offset):
        '''
        字符串类型
        :param new_offset:
        :return:
        '''
        data_len = int_from_byte(self.byte_data[new_offset])
        new_offset += 1
        temp_offset = new_offset + data_len
        return self.byte_data[new_offset:temp_offset].decode('utf-8'), temp_offset

    def decode_long_str(self, new_offset):
        '''
        长文本类型
        :param new_offset:
        :return:
        '''
        size = int_from_byte(self.byte_data[new_offset])
        new_offset += 1
        data_len = uint_from_bytes(self.byte_data[new_offset:new_offset + size])
        new_offset += size
        temp_offset = new_offset + data_len
        return self.byte_data[new_offset:temp_offset].decode('utf-8'), temp_offset

    def decode_uint(self, new_offset):
        '''
        无符号int类型
        :param new_offset:
        :return:
        '''
        data_len = int_from_byte(self.byte_data[new_offset])
        new_offset += 1
        temp_offset = new_offset + data_len
        value = uint_from_bytes(self.byte_data[new_offset:temp_offset])
        return value, temp_offset

    def decode_int(self, new_offset):
        '''
        有符号int
        :param new_offset:
        :return:
        '''
        temp_offset = new_offset + 4
        value = struct.unpack(b'!i', self.byte_data[new_offset:temp_offset])[0]
        return value, temp_offset

    def decode_float(self, new_offset):
        '''
        float类型
        :param new_offset:
        :return:
        '''
        temp_offset = new_offset + 4
        value = struct.unpack(b'!f', self.byte_data[new_offset:temp_offset])[0]
        return value, temp_offset

    def decode_double(self, new_offset):
        '''
        doubele类型
        :param new_offset:
        :return:
        '''
        temp_offset = new_offset + 8
        value = struct.unpack(b'!d', self.byte_data[new_offset:temp_offset])[0]
        return value, temp_offset

    def decode_pointer(self, offset):
        '''
        指针
        '''
        data_len = int_from_byte(self.byte_data[offset])  # 数据范围 1-5字节
        offset += 1
        new_offset = offset + data_len
        buf = uint_from_bytes(self.byte_data[offset:new_offset])  # \x0\x0\x0\I
        pointer = self.base_offset + buf
        value, _ = self.decode_content(pointer)
        return value, new_offset

    decode_data_type = {
        1: decode_arrary,
        2: decode_pointer,
        3: decode_str,
        4: decode_long_str,
        5: decode_uint,
        6: decode_int,
        7: decode_float,
        8: decode_double,
    }

    def decode_content(self, offset):
        data_type = int_from_byte(self.byte_data[offset])
        new_offset = offset + 1
        if data_type in self.decode_data_type.keys():
            decoder_func = self.decode_data_type[data_type]
        else:
            raise ValueError('Unexpected data type number ({type}) ,please check!'.format(type=data_type))
        return decoder_func(self, new_offset)

    def decode_content_structure(self, offset):
        content, _ = self.decode_content(offset)
        content = map_key_value(self.columns, content)
        return content

    def decode_content_direct(self, offset):
        data_len = uint_from_bytes(self.byte_data[offset:offset + 4])
        offset += 4
        temp_list = self.byte_data[offset:offset + data_len].decode('utf-8').split("\t")
        dic = {}
        for i, value in enumerate(self.columns):
            dic[value] = temp_list[i]
        return dic

    parse_type_func = {
        1: decode_content_structure,
        2: decode_content_direct
    }

    def read_awdb(self, ip_address):
        '''
        读取awdb
        '''
        node_index = self.find_tree_index(ip_address)
        if node_index == '':
            return {}
        pointer = self.base_offset + node_index - self.node_count - 10
        parse_func = self.parse_type_func[self.decode_type]
        content = parse_func(self, pointer)
        return content
